from pathlib import Path
import docx
# from docx.document import document
from typing import Dict, Any
from docx.oxml import OxmlElement
from docx.oxml.ns import qn
from docx.shared import RGBColor
from docx.oxml.shared import qn
from datetime import datetime
from lxml import etree
import logging
import os
import zipfile
import xml.dom.minidom
import requests
import json
import sys
from os.path import dirname, abspath
import shutil
sys.path.append(dirname(dirname(abspath(__file__))))
from utils.models import LLMClient, ApiClient
from utils.models import DocumentEvaluation
class DocumentProcessor:
    def __init__(self):
        self.grading_rules = {
            'font_size': 12,
            'font_name': 'Times New Roman',
            'line_spacing': 1.5,
            'margins': {
                'top': 2.54,
                'bottom': 2.54,
                'left': 3.18,
                'right': 3.18
            }
        }
        self.logger = None
        self.pages = None
        self.client = None
        self.llm_type = "deepseek"
        self.required_sections = None

    def set_logger(self, logger):
        """设置日志记录器"""
        self.logger = logger
        # 同时为客户端设置logger
        self.llm_client = LLMClient(logger=logger)
        # self.api_client = ApiClient(
        #     api_key="sk-a2200d958f4c4bc8938e77fa97edfd21",  # 请确保这是有效的API密钥
        #     base_url="https://api.deepseek.com/v1",  # 使用正确的API版本端点
        #     logger=logger
        # )

    def grade_document(self, doc_path: Path, student_id: str, student_name: str, required_sections: list) -> Dict[str, Any]:
        """
        Grade a document and return the results
        """
        # 获取文档名称（不包含扩展名）
        doc_name = doc_path.stem
        
        try:
            doc = docx.Document(doc_path)
            
            # 记录开始评分信息
            # if self.logger:
            #     self.logger.info(f"\n开始评分文档: {doc_name}")
            #     self.logger.info(f"学号: {student_id}")
            #     self.logger.info(f"姓名: {student_name}")

            self.document_to_pages(doc_path)
            # 打印文档统计信息
            doc_stats = self.statistics_document(doc, doc_path)
            
            # 检查文档结构完整性
            self.required_sections = required_sections
            structure_completeness = self.check_structure_completeness(doc, required_sections)
            
            # 输出结构完整性检查结果
            if self.logger:
                self.logger.info("\n文档结构完整性检查:")
                for section, exists in structure_completeness.items():
                    status = "✓ 存在" if exists else "✗ 缺失"
                    self.logger.info(f"{section}: {status}")
            
            # # 检查格式
            # format_check = self._check_formatting(doc)
            
                        # 最后评估文档排版
            # if self.logger:
            #     self.logger.info("\n============评估文档排版============")
            evaluation = self.output_pages()
            
            # 获取所有日志内容
            log_content = ""
            if self.logger and hasattr(self.logger, 'handlers'):
                for handler in self.logger.handlers:
                    if isinstance(handler, logging.FileHandler):
                        log_file = handler.baseFilename
                        with open(log_file, 'r', encoding='utf-8') as f:
                            log_content = f.read()
            
            # 从配置文件读取system_prompt
            system_prompt = None
            try:
                import json
                import os
                config_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'config.json')
                with open(config_path, 'r', encoding='utf-8') as f:
                    config = json.load(f)
                    system_prompt = config.get('system_prompt')
            except Exception as e:
                if self.logger:
                    self.logger.warning(f"无法读取配置文件中的system_prompt: {str(e)}")
            
            # 使用文档内容进行评估
            # 将文档内容转换为纯文本
            doc_text = ""
            for para in doc.paragraphs:
                if para.text.strip():
                    doc_text += para.text.strip() + "\n"
            
            # 使用文档内容进行评估
            evaluation = self.evaluate(doc_text, system_prompt)
        
            
            # 构建批注内容
            comment_text = f"""
                总分: {evaluation.score}
                评语: {evaluation.comments}
            """
            
            # 返回包含评分信息的字典
            return {
                'comment_text': comment_text,
                'score': evaluation.score,
                'comments': evaluation.comments,
                'student_id': student_id,
                'student_name': student_name
            }
            
        except Exception as e:
            if self.logger:
                self.logger.error(f"评分过程中发生错误: {str(e)}")
            raise

    def _check_formatting(self, doc: docx.Document) -> Dict[str, Any]:
        """
        Check if images are centered and have centered captions below them
        """
        total_images = 0
        centered_images = 0
        centered_captions = 0
        non_centered_images = 0
        non_centered_captions = 0
        missing_captions = 0
        image_groups = []  # 用于存储图片组
        current_group = []  # 当前正在处理的图片组
        image_captions = []  # 存储图片说明

        try:
            i = 0
            while i < len(doc.paragraphs):
                para = doc.paragraphs[i]
                if not para or not hasattr(para, 'runs') or not para.runs:
                    i += 1
                    continue
                    
                # 检查当前段落是否包含图片
                has_image = False
                for run in para.runs:
                    if (run is not None and 
                        hasattr(run, 'element') and 
                        run.element is not None and 
                        hasattr(run.element, 'xml') and 
                        run.element.xml is not None):
                        try:
                            xml_str = str(run.element.xml)
                            if len(xml_str) > 0 and 'w:drawing' in xml_str:
                                has_image = True
                                total_images += 1
                                
                                # 检查段落是否居中
                                if self._is_paragraph_centered(para):
                                    centered_images += 1
                                else:
                                    non_centered_images += 1

                                # 将图片信息添加到当前组
                                current_group.append(total_images)
                                
                        except Exception as e:
                            continue

                # 如果当前段落包含图片，检查下一段是否也是图片
                if has_image:
                    next_para = doc.paragraphs[i + 1] if i + 1 < len(doc.paragraphs) else None
                    if next_para is not None and hasattr(next_para, 'runs') and len(next_para.runs) > 0:
                        has_next_image = False
                        for run in next_para.runs:
                            if (run is not None and 
                                hasattr(run, 'element') and 
                                run.element is not None and 
                                hasattr(run.element, 'xml') and 
                                run.element.xml is not None):
                                try:
                                    xml_str = str(run.element.xml)
                                    if len(xml_str) > 0 and 'w:drawing' in xml_str:
                                        has_next_image = True
                                        break
                                except:
                                    pass
                        
                        if has_next_image:
                            i += 1
                            continue

                # 如果当前段落没有图片，或者下一段不是图片，则处理当前组的说明
                if current_group:
                    # 检查说明
                    caption_para = doc.paragraphs[i + 1] if i + 1 < len(doc.paragraphs) else None
                    caption_found = False
                    caption_text = ""
                    
                    if caption_para and hasattr(caption_para, 'text') and caption_para.text:
                        caption_text = caption_para.text.strip()
                        if caption_text:
                            # 检查是否是图片标题
                            is_caption = True  # 默认认为是图片标题
                            
                            # 检查段落样式是否为分级标题
                            if hasattr(caption_para, 'style') and caption_para.style:
                                style_name = caption_para.style.name.lower()
                                # 如果是分级标题（如"标题1"、"标题2"、"标题3"等），则不是图片标题
                                if style_name in ['标题1', '标题2', '标题3', '标题4', '标题5', '标题6', 
                                                'heading 1', 'heading 2', 'heading 3', 'heading 4', 'heading 5', 'heading 6']:
                                    is_caption = False
                            
                            if is_caption:
                                caption_found = True
                                # 检查段落是否居中
                                if self._is_paragraph_centered(caption_para):
                                    centered_captions += 1
                                else:
                                    non_centered_captions += 1

                    if not caption_found:
                        missing_captions += 1
                        caption_text = "（无标题）"

                    # 记录图片说明
                    for img_num in current_group:
                        image_captions.append(f"图片 #{img_num} 标题: {caption_text}")

                    # 保存当前组并清空
                    image_groups.append(current_group)
                    current_group = []
                    i += 2  # 跳过说明段落
                else:
                    i += 1

            # 输出统计信息到日志
            if self.logger:
                self.logger.info("\n图片格式检查统计:")
                self.logger.info(f"总图片数: {total_images}")
                self.logger.info(f"图片组数: {len(image_groups)}")
                self.logger.info(f"居中图片数: {centered_images}")
                self.logger.info(f"未居中图片数: {non_centered_images}")
                self.logger.info(f"居中标题数: {centered_captions}")
                self.logger.info(f"未居中标题数: {non_centered_captions}")
                self.logger.info(f"缺失标题数: {missing_captions}")

                # 输出所有图片说明，按组输出
                self.logger.info("\n图片列表:")
                for group_idx, group in enumerate(image_groups, 1):
                    # 获取该组的第一个图片说明作为图表题
                    first_caption = next((c for c in image_captions if f"图片 #{group[0]}" in c), None)
                    caption_text = first_caption.split("标题: ")[1] if first_caption else "（无标题）"
                    self.logger.info(f"\n图片组 #{group_idx}，包含{len(group)}个图片，图表题是：{caption_text}")

        except Exception as e:
            if self.logger:
                self.logger.error(f"格式检查过程中发生错误: {str(e)}")

        return {
            'image_formatting': {
                'status': 'pass' if non_centered_images == 0 and non_centered_captions == 0 and missing_captions == 0 else 'fail',
                'statistics': {
                    'total_images': total_images,
                    'image_groups': len(image_groups),
                    'centered_images': centered_images,
                    'non_centered_images': non_centered_images,
                    'centered_captions': centered_captions,
                    'non_centered_captions': non_centered_captions,
                    'missing_captions': missing_captions
                }
            }
        }

    def _check_content(self, doc: docx.Document) -> Dict[str, Any]:
        """
        Check document content
        """
        content_results = {
            'has_table_of_contents': self._check_table_of_contents(doc),
            'content_length': self._check_content_length(doc),
            'references': self._check_references(doc)
        }
        return content_results

    def _check_structure(self, doc: docx.Document) -> Dict[str, Any]:
        """
        Check document structure
        """
        structure_results = {
            'heading_structure': self._check_heading_structure(doc),
            'paragraph_structure': self._check_paragraph_structure(doc)
        }
        return structure_results

    def _check_fonts(self, doc: docx.Document) -> Dict[str, Any]:
        """
        Check font consistency
        """
        expected_font_name = self.grading_rules['font_name']
        expected_font_size = self.grading_rules['font_size']
        font_issues = []

        for para in doc.paragraphs:
            for run in para.runs:
                if run.font.name != expected_font_name:
                    font_issues.append(f"Font name mismatch in paragraph: {para.text}")
                if run.font.size and run.font.size.pt != expected_font_size:
                    font_issues.append(f"Font size mismatch in paragraph: {para.text}")

        return {
            'status': 'pass' if not font_issues else 'fail',
            'details': font_issues
        }

    def _create_comment_element(self, text: str) -> OxmlElement:
        """
        创建批注元素
        """
        comment = OxmlElement('w:comment')
        comment.set(qn('w:id'), '0')
        comment.set(qn('w:author'), '教师')
        comment.set(qn('w:date'), '2024-05-09T00:00:00Z')
        comment.set(qn('w:initials'), '教师')
        
        # 创建批注文本
        text_element = OxmlElement('w:p')
        run = OxmlElement('w:r')
        text_run = OxmlElement('w:t')
        text_run.text = text
        run.append(text_run)
        text_element.append(run)
        comment.append(text_element)
        
        return comment

    def _add_comment_to_paragraph(self, paragraph, comment_text: str):
        """
        向段落添加批注
        """
        # 创建批注元素
        comment = self._create_comment_element(comment_text)
        
        # 获取文档的批注部分
        comments_part = paragraph.part.get_or_add_part('comments')
        comments_element = comments_part.element
        
        # 添加批注
        comments_element.append(comment)
        
        # 在段落中添加批注引用
        comment_reference = OxmlElement('w:commentReference')
        comment_reference.set(qn('w:id'), '0')
        paragraph._p.append(comment_reference)

    def save_results(self, results: str, output_path: Path, original_path: Path, author: str) -> None:
        """
        Save grading results by adding comments to the first line of the generated document
        """
        # 复制原始文档到输出路径
        shutil.copy2(original_path, output_path)

        # 加载生成的文档
        doc = docx.Document(output_path)

        # 在文档的第一句话添加批注
        first_para = None
        for para in doc.paragraphs:
            if para.text.strip():
                first_para = para
                break

        if first_para is not None:
            # 使用bayoo-docx的批注功能
            first_para.add_comment(results, author=author)

        # 保存到输出路径
        doc.save(output_path)

    def check_structure_completeness(self, doc: docx.Document, required_sections: list) -> Dict[str, bool]:
        """
        检查文档结构是否完整，包含配置文件中指定的部分
        同时输出每个部分的页数和文档总页数
        """
        found_sections = {section: False for section in required_sections}
        section_positions = []  # 存储每个部分的位置
        self.required_sections = required_sections
        # 遍历所有页，查找每个部分
        for section in required_sections:
            for page_number, page_content in enumerate(self.pages, start=1):
                if section in page_content:
                    found_sections[section] = True
                    section_positions.append((section, page_number))
                    break

        # 计算文档总页数
        total_pages = len(self.pages)
        if self.logger:
            self.logger.info(f"\n文档总页数: {total_pages}")

        # 计算每个部分的页数
        for i, (section, start_page) in enumerate(section_positions):
            end_page = section_positions[i + 1][1] - 1 if i + 1 < len(section_positions) else total_pages
            pages = end_page - start_page + 1
            if self.logger:
                self.logger.info(f"'{section}' 的页数: 第{start_page}页 - 第{end_page}页，共{pages}页")

        return found_sections

    def _is_paragraph_centered(self, paragraph) -> bool:
        """
        判断段落是否居中
        首先检查样式的对齐方式，如果没有样式或样式没有设置对齐方式，则检查段落的对齐方式
        """
        # 首先检查样式的对齐方式
        if hasattr(paragraph, 'style') and paragraph.style:
            if hasattr(paragraph.style, 'paragraph_format') and paragraph.style.paragraph_format:
                if paragraph.style.paragraph_format.alignment == docx.enum.text.WD_ALIGN_PARAGRAPH.CENTER:
                    return True
        
        # 如果样式没有设置对齐方式，则检查段落的对齐方式
        if hasattr(paragraph, 'paragraph_format') and paragraph.paragraph_format:
            if paragraph.paragraph_format.alignment == docx.enum.text.WD_ALIGN_PARAGRAPH.CENTER:
                return True
        
        return False

    def statistics_document(self, doc: docx.Document, doc_path: Path = None) -> Dict[str, int]:
        """
        统计文档的基本信息
        返回：总页数、总字符数、图片数量
        """
        if self.logger:
            self.logger.info("\n开始统计文档信息...")
            
        # 统计总字符数
        total_chars = 0
        for para in doc.paragraphs:
            total_chars += len(para.text)

        # 统计图片数量
        total_images = 0
        for para in doc.paragraphs:
            # if self.logger:
            #     self.logger.info(f"检查段落是否包含图片: {para.text[:50]}...")
                
            for run in para.runs:
                if (run is not None and 
                    hasattr(run, 'element') and 
                    run.element is not None and 
                    hasattr(run.element, 'xml') and 
                    run.element.xml is not None):
                    try:
                        xml_str = str(run.element.xml)
                        if len(xml_str) > 0 and 'w:drawing' in xml_str:
                            total_images += 1
                            # if self.logger:
                            #     self.logger.info(f"找到第 {total_images} 个图片")
                    except Exception as e:
                        if self.logger:
                            self.logger.warning(f"处理XML时发生错误: {str(e)}")
                        continue

        # 获取文档的实际页数
        total_pages = 1
        try:
            if doc_path:
                # if self.logger:
                #     self.logger.info("从文档属性中获取页数...")
                # 打开docx文件（实际上是一个zip文件）
                with zipfile.ZipFile(doc_path) as document:
                    # 读取docProps/app.xml文件
                    dxml = document.read('docProps/app.xml')
                    # 解析XML内容以获取Pages节点中的数值
                    xml_content = dxml.decode('utf-8')
                    import re
                    pages_match = re.search(r'<Pages>(\d+)</Pages>', xml_content)
                    if pages_match:
                        total_pages = int(pages_match.group(1))
                        # if self.logger:
                        #     self.logger.info(f"从XML中获取到页数: {total_pages}")
                    # else:
                    #     if self.logger:
                    #         self.logger.warning("未能在XML中找到页数信息")
        except Exception as e:
            if self.logger:
                self.logger.warning(f"从XML获取页数时发生错误: {str(e)}")
            # 如果获取失败，使用简单的方法
            try:
                if hasattr(doc, '_element') and doc._element is not None:
                    page_breaks = doc._element.xpath('.//w:br[@w:type="page"]')
                    section_breaks = doc._element.xpath('.//w:sectPr')
                    total_pages = len(page_breaks) + len(section_breaks) + 1
                    if self.logger:
                        self.logger.info(f"通过分页符计算得到页数: {total_pages}")
            except Exception as e:
                if self.logger:
                    self.logger.warning(f"无法获取实际页数: {str(e)}")

        if self.logger:
            self.logger.info("\n文档统计信息:")
            self.logger.info(f"总页数: {total_pages}")
            self.logger.info(f"总字符数: {total_chars}")
            self.logger.info(f"图片数量: {total_images}")

        return {
            'total_pages': total_pages,
            'total_chars': total_chars,
            'total_images': total_images
        }

    def document_to_pages(self, doc_path: Path) -> None:
        """
        将Word文档按页分割，并打印结果,删除其中的图片要素
        
        Args:
            doc_path: Word文档路径
        """
        try:
            # 打开docx文件（实际上是一个zip文件）
            with zipfile.ZipFile(doc_path, 'r') as doc:
                # 读取document.xml
                xml_content = doc.read('word/document.xml')
                
                # 解析XML
                tree = etree.fromstring(xml_content)
                
                # 定义命名空间
                namespaces = {
                    'w': 'http://schemas.openxmlformats.org/wordprocessingml/2006/main',
                    'a': 'http://schemas.openxmlformats.org/drawingml/2006/main'
                }
                
                # 查找所有a:graphic和a:graphicData元素
                graphic_elements = tree.xpath('//a:graphic', namespaces=namespaces)
                graphic_data_elements = tree.xpath('//a:graphicData', namespaces=namespaces)
                
                # 删除找到的元素
                for element in graphic_elements + graphic_data_elements:
                    parent = element.getparent()
                    if parent is not None:
                        parent.remove(element)
                
                # 将修改后的XML转换为字符串
                modified_xml = etree.tostring(tree, encoding='UTF-8').decode('utf-8')
                
                # 使用正则表达式分割字符串
                import re
                self.pages = re.split(r'<w:lastRenderedPageBreak/>', modified_xml, flags=re.IGNORECASE)

        except Exception as e:
            if self.logger:
                self.logger.error(f"删除图形数据时发生错误: {str(e)}")
            raise

    def output_pages(self) -> None:
        """
        将页面内容输出到日志
        """
        if self.logger and self.pages:
            self.logger.info("\n============以下是word文档的xml格式内容============")
            for i, page in enumerate(self.pages, 1):
                self.logger.info(f"\n--- 第{i}页 ---")
                self.logger.info(page)
            self.logger.info("==============================\n")

    def evaluate(self, doc_content: str, system_prompt: str = None) -> DocumentEvaluation:
        """
        将文档内容发送给AI模型进行评估
        
        Args:
            doc_content: 文档的纯文本内容
            system_prompt: 系统提示词（可选）
            
        Returns:
            DocumentEvaluation对象，包含评分和评语
        """
        if self.llm_type == "deepseek":
            from utils.models import ApiClient
            self.client = ApiClient(
                logger=self.logger
            )
        elif self.llm_type == "ollama":
            from utils.models import LLMClient
            self.client = LLMClient(
                logger=self.logger
            )
        # 使用AI模型进行评估
        evaluation = self.client.evaluate_document(doc_content, self.required_sections, system_prompt)
        
        if self.logger:
            self.logger.info("\n=============== 文档评估结果 ===============")
            self.logger.info(f"总分: {evaluation.score}")
            self.logger.info("\n总体评语:")
            self.logger.info(evaluation.comments)
            self.logger.info("==========================================\n")
            
        return evaluation